"""
ConnLLM图像输入能力
定义处理图像输入的接口和基本实现
"""
import base64
import os
import mimetypes
import logging
from abc import abstractmethod
from typing import Dict, Any, List, Optional, Union
import requests
from urllib.parse import urlparse

from .base import Capability, register_capability

logger = logging.getLogger("connllm.capabilities.image_input")

@register_capability("image_input")
class ImageInputCapability(Capability):
	"""图像输入能力基类，处理图像URL和二进制数据"""
	
	def __init__(self, provider):
		"""初始化能力"""
		super().__init__(provider)
		self._supported_formats = ["image/jpeg", "image/png", "image/webp"]
	
	def is_supported(self) -> bool:
		"""检查提供商是否支持图像输入能力"""
		# 由具体实现重写
		return False
		
	def create_image_message(self, 
							text: str, 
							image_source: Union[str, bytes], 
							role: str = "user", 
							detail_level: str = "auto") -> Dict[str, Any]:
		"""
		创建包含图像的消息
		
		Args:
			text: 文本内容
			image_source: 图像URL或二进制数据
			role: 消息角色，默认为user
			detail_level: 图像细节级别 (auto/high/low)
			
		Returns:
			包含图像的消息字典
		"""
		# 处理图像来源
		image_data = self._process_image_source(image_source)
		
		# 由子类实现具体消息格式
		return self._format_image_message(text, image_data, role, detail_level)
	
	def _process_image_source(self, image_source: Union[str, bytes]) -> Dict[str, Any]:
		"""
		处理图像来源，返回统一格式
		
		Args:
			image_source: 图像URL或二进制数据
			
		Returns:
			处理后的图像数据字典
		"""
		# 如果是字符串，检查是否为URL或文件路径
		if isinstance(image_source, str):
			# 检查是否为URL
			parsed = urlparse(image_source)
			if parsed.scheme in ['http', 'https']:
				return {
					"type": "url",
					"url": image_source,
					"mimetype": None  # URL不需要指定mimetype
				}
			
			# 否则视为文件路径
			if os.path.isfile(image_source):
				with open(image_source, 'rb') as f:
					binary_data = f.read()
				
				# 获取mimetype
				mime_type = mimetypes.guess_type(image_source)[0]
				if not mime_type or mime_type not in self._supported_formats:
					mime_type = "image/jpeg"  # 默认JPEG
					
				# 验证图像大小
				self._validate_image_size(binary_data)
					
				return {
					"type": "binary",
					"data": binary_data,
					"mimetype": mime_type
				}
			else:
				raise ValueError(f"图像文件不存在: {image_source}")
		
		# 如果是二进制数据
		elif isinstance(image_source, bytes):
			# 尝试检测图像类型
			mime_type = self._detect_image_type(image_source)
			
			# 验证图像大小
			self._validate_image_size(image_source)
			
			return {
				"type": "binary",
				"data": image_source,
				"mimetype": mime_type
			}
		
		else:
			raise TypeError("图像来源必须是URL字符串、文件路径或二进制数据")
	
	def _fetch_image_from_url(self, url: str) -> bytes:
		"""
		从URL获取图像数据，增强错误处理
		
		Args:
			url: 图像URL
			
		Returns:
			图像二进制数据
		"""
		try:
			response = requests.get(url, timeout=10)
			response.raise_for_status()
			return response.content
		except requests.exceptions.Timeout:
			raise ValueError(f"获取图像超时: {url}")
		except requests.exceptions.HTTPError as e:
			raise ValueError(f"获取图像失败，HTTP错误({e.response.status_code}): {url}")
		except requests.exceptions.RequestException as e:
			raise ValueError(f"获取图像失败: {str(e)}")
	
	def _detect_image_type(self, data: bytes) -> str:
		"""
		检测图像类型
		
		Args:
			data: 图像二进制数据
			
		Returns:
			MIME类型
		"""
		# 简单的图像类型检测
		if data.startswith(b'\xff\xd8\xff'):
			return "image/jpeg"
		elif data.startswith(b'\x89PNG\r\n\x1a\n'):
			return "image/png"
		elif data.startswith(b'RIFF') and data[8:12] == b'WEBP':
			return "image/webp"
		else:
			# 默认JPEG
			return "image/jpeg"
	
	def _validate_image_size(self, data: bytes, max_size_mb: float = 20.0) -> None:
		"""
		检查图像大小是否在合理范围内
		
		Args:
			data: 图像二进制数据
			max_size_mb: 最大大小(MB)
		"""
		size_mb = len(data) / (1024 * 1024)
		if size_mb > max_size_mb:
			logger.warning(f"图像大小({size_mb:.2f}MB)超过建议值({max_size_mb}MB)")
	
	def _encode_image_base64(self, data: bytes, mime_type: str) -> str:
		"""
		将图像编码为base64
		
		Args:
			data: 图像二进制数据
			mime_type: MIME类型
			
		Returns:
			base64编码的图像数据URL
		"""
		b64_data = base64.b64encode(data).decode('utf-8')
		return f"data:{mime_type};base64,{b64_data}"
	
	@abstractmethod
	def _format_image_message(self, 
							 text: str, 
							 image_data: Dict[str, Any], 
							 role: str, 
							 detail_level: str) -> Dict[str, Any]:
		"""
		格式化图像消息，由具体提供商实现
		
		Args:
			text: 文本内容
			image_data: 图像数据
			role: 消息角色
			detail_level: 图像细节级别
			
		Returns:
			格式化的消息字典
		"""
		pass
	
	@classmethod
	def create_for_provider(cls, provider):
		"""
		根据提供商创建合适的图像输入能力实现
		
		Args:
			provider: 提供商实例
			
		Returns:
			图像输入能力实现实例
		"""
		provider_type = provider.config_manager.get_provider_type()
		
		if provider_type == "openrouter":
			from ..implementations.openrouter.image_input import OpenRouterImageInput
			return OpenRouterImageInput(provider)
			
		# 未来可以添加更多提供商支持
		
		raise NotImplementedError(f"图像输入能力不支持提供商: {provider_type}")
